"""6.4 【上下文类】Context 消息回调类MCP开发"""
import asyncio
from contextlib import AsyncExitStack
from typing import Optional
from mcp import ClientSession
from mcp.client.sse import sse_client

class MCPClient:

    def __init__(self):
        #创建线程管理栈
        self.exit_stack = AsyncExitStack()
        self.session: Optional[ClientSession] = None

    async def logging_handler(self,params):
        print(f"接收到服务日志消息: {params.data}")

    """创建连接服务端"""
    async def connect_to_server(self, server_path: str):

        # 一、创建sse_client
        client = sse_client(url=server_path)
        stdio_transport = await self.exit_stack.enter_async_context(client)
        read_stream, write_stream = stdio_transport

        # 三、创建ClientSession,日志回调函数logging_callback，实时接收服务端日志状态信息
        client_session = ClientSession(
            read_stream,
            write_stream,
            logging_callback=self.logging_handler,
        )
        self.session = await self.exit_stack.enter_async_context(client_session)
        # 四、初始化session
        await self.session.initialize()
        #五、列出所有可用的工具
        response = await self.session.list_tools()
        tools = response.tools
        tool = tools[0]
        print("\nConnected to server with tools:",tools)
        # 调用工具链
        result = await self.session.call_tool(name=tool.name,arguments={"batch_list":["one","two","three"]})
        print(result)


    #关闭连接
    async def cleanup(self):
        await self.exit_stack.aclose()

async def main():
    client = MCPClient()
    try:
        await client.connect_to_server("http://localhost:8000/sse")

    except Exception as e:
        print(f"Error: {str(e)}")
    finally:
        await client.cleanup()

if __name__ == "__main__":
    asyncio.run(main())